import gradio as gr
import os
import json
from openai import OpenAI
import uvicorn
from fastapi import FastAPI, HTTPException, UploadFile, File
from pydantic import BaseModel
from typing import Optional, List, Dict
from io import BytesIO
from docx import Document
from PyPDF2 import PdfReader
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# --- 配置部分 ---
# 从环境变量获取百度AI API Key
BAIDU_API_KEY = os.environ.get("BAIDU_API_KEY", "xxx")  # 百度API Key
BAIDU_BASE_URL = "https://aistudio.baidu.com/llm/lmapi/v3"
MODEL_NAME = "ernie-4.5-turbo-vl"  # 百度模型

# 初始化OpenAI客户端以兼容百度API
client = OpenAI(
    api_key=BAIDU_API_KEY,
    base_url=BAIDU_BASE_URL,
)

# 支持的合同类型、对应的审核维度和常见风险点
CONTRACT_CONFIG = {
    "劳动合同": {
        "dimensions": [
            "条款完整性",
            "法律合规性",
            "用工风险",
            "薪酬福利条款",
            "解约条款",
        ],
        "focus_points": [
            "试用期约定",
            "薪资结构",
            "竞业限制",
            "社会保险",
            "解除合同条款",
        ],
    },
    "租赁合同": {
        "dimensions": [
            "条款完整性",
            "租赁期限合规",
            "租金支付条款",
            "维修责任划分",
            "违约责任",
        ],
        "focus_points": ["租赁期限", "租金支付", "违约责任", "房屋修缮", "转租条款"],
    },
    "销售合同": {
        "dimensions": [
            "产品描述准确性",
            "价格与支付条款",
            "交货与验收标准",
            "质量保证条款",
            "违约责任",
        ],
        "focus_points": ["价款支付", "交货方式", "质量标准", "验收条款", "违约责任"],
    },
    "服务合同": {
        "dimensions": [
            "服务范围明确性",
            "服务标准量化",
            "付款条件合理性",
            "服务期限明确性",
            "保密条款",
        ],
        "focus_points": ["服务范围", "验收标准", "付款条件", "保密条款", "违约责任"],
    },
    "保密协议": {
        "dimensions": [
            "保密范围界定",
            "保密期限合理性",
            "泄密责任明确性",
            "例外情况条款",
            "知识产权保护",
        ],
        "focus_points": ["保密范围", "保密期限", "违约责任", "例外情况", "知识产权"],
    },
    "合作协议": {
        "dimensions": [
            "合作范围明确性",
            "权利义务划分",
            "出资与收益分配",
            "决策机制设置",
            "退出机制",
        ],
        "focus_points": ["出资比例", "利润分配", "决策权分配", "退出机制", "保密条款"],
    },
    "采购合同": {
        "dimensions": [
            "货物规格明确性",
            "验收标准量化",
            "付款条件合理性",
            "交付时间约定",
            "违约责任",
        ],
        "focus_points": ["货物规格", "验收标准", "付款条件", "违约责任", "争议解决"],
    },
    "借款合同": {
        "dimensions": [
            "借款金额与用途",
            "利率合规性",
            "还款方式明确性",
            "担保条款有效性",
            "逾期责任",
        ],
        "focus_points": ["借款金额", "利率约定", "还款方式", "担保条款", "违约责任"],
    },
    "股权转让协议": {
        "dimensions": [
            "转让标的明确性",
            "价格与支付条款",
            "交割条件设置",
            "陈述与保证条款",
            "税费承担",
        ],
        "focus_points": ["转让价格", "支付方式", "股权交割", "陈述与保证", "税费承担"],
    },
    "知识产权许可协议": {
        "dimensions": [
            "许可范围界定",
            "许可期限合理性",
            "使用费支付条款",
            "侵权责任划分",
            "保密条款",
        ],
        "focus_points": ["许可范围", "许可期限", "使用费", "侵权责任", "保密条款"],
    },
}

# 风险等级说明
RISK_LEVELS = {
    "高": "可能导致重大法律责任或经济损失，必须修改",
    "中": "存在一定风险，建议修改以避免潜在问题",
    "低": "风险较小，但仍有优化空间",
}


# --- 文档处理工具函数 ---
def extract_text_from_file(file: UploadFile) -> str:
    """从上传的文件中提取文本内容"""
    try:
        content = file.file.read()
        file_ext = os.path.splitext(file.filename)[1].lower()

        if file_ext == ".txt":
            return content.decode("utf-8", errors="ignore")

        elif file_ext == ".docx":
            doc = Document(BytesIO(content))
            return "\n".join([para.text for para in doc.paragraphs])

        elif file_ext == ".pdf":
            reader = PdfReader(BytesIO(content))
            text = ""
            for page in reader.pages:
                extracted_text = page.extract_text()
                if extracted_text:
                    text += extracted_text + "\n"
            return text

        else:
            return f"不支持的文件格式: {file_ext}。请上传TXT、DOCX或PDF文件。"

    except Exception as e:
        return f"文件处理错误: {str(e)}"


# --- 核心处理函数 ---
def review_contract_core(
    contract_type: str,
    contract_text: str,
    review_dimensions: List[str],
    focus_points: str = "",
) -> dict:
    """合同审核核心逻辑"""
    try:
        # 构建提示词
        prompt = f"""
        你是一位拥有10年以上经验的合同法律师，擅长各类合同的风险审核与条款优化。
        
        请根据以下信息，对合同进行专业审核：
        - 合同类型：{contract_type}
        - 审核维度：{', '.join(review_dimensions)}
        - 特别关注点：{focus_points if focus_points else '无特别关注点'}
        
        请按照以下结构提供审核结果：
        1. 风险评估摘要：总体风险等级（高/中/低）及简要说明
        2. 风险点识别：逐条列出合同中存在的风险点，每条包含：
           - 条款位置：风险条款在合同中的大致位置
           - 风险描述：具体风险内容
           - 风险等级：高/中/低
           - 法律依据：相关法律法规依据（如有）
        3. 条款优化建议：对有问题的条款提供具体修改建议，格式为"原条款建议：XXX → 修改后：XXX"
        4. 缺失条款提示：根据合同类型，提示应包含但缺失的重要条款
        5. 审核结论：总结性意见和签署前注意事项
        
        请基于提供的合同文本进行分析，语言专业准确，建议具有可操作性。
        对于识别的风险点，请优先考虑与当前合同类型最相关的问题。
        """

        # 调用百度AI API
        messages = [
            {
                "role": "system",
                "content": "你是一位专业的合同法律师，精通各类合同审核和风险评估。",
            },
            {"role": "user", "content": f"合同文本：{contract_text}\n\n{prompt}"},
        ]

        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            stream=False,
            temperature=0.3,  # 降低创造性，提高审核准确性
        )

        ai_response = completion.choices[0].message.content
        return {"success": True, "result": ai_response}

    except Exception as e:
        print(f"调用AI API时出错: {e}")
        return {"success": False, "error": f"审核失败：{str(e)}"}


# --- Gradio界面函数 ---
def gradio_interface():
    with gr.Blocks(title="合同智审 - AI合同审核专家") as demo:
        gr.Markdown("## 📜 合同智审 - AI合同审核专家")
        gr.Markdown("上传合同文件或输入合同文本，获取专业的合同风险审核与优化建议")

        with gr.Row():
            with gr.Column(scale=1):
                contract_type = gr.Dropdown(
                    label="合同类型",
                    choices=list(CONTRACT_CONFIG.keys()),
                    value="劳动合同",
                )

                review_dimensions = gr.CheckboxGroup(
                    label="审核维度",
                    choices=CONTRACT_CONFIG["劳动合同"]["dimensions"],
                    value=CONTRACT_CONFIG["劳动合同"]["dimensions"],
                )

                focus_points = gr.Textbox(
                    label="特别关注点（可选）",
                    placeholder=f"推荐关注点：{', '.join(CONTRACT_CONFIG['劳动合同']['focus_points'])}",
                )

                file_upload = gr.File(
                    label="上传合同文件（可选）",
                    file_types=[".txt", ".docx", ".pdf"],
                    height=100,
                )

                contract_text = gr.Textbox(
                    label="或直接输入合同文本",
                    placeholder="请粘贴合同文本内容...",
                    lines=8,
                )

                analyze_btn = gr.Button("开始审核", variant="primary")

            with gr.Column(scale=2):
                result = gr.Markdown(label="合同审核结果")

        # 当合同类型变化时，更新审核维度和推荐关注点
        def update_contract_config(contract_type):
            config = CONTRACT_CONFIG[contract_type]
            return [
                gr.CheckboxGroup(
                    choices=config["dimensions"], value=config["dimensions"]
                ),
                gr.Textbox(
                    placeholder=f"推荐关注点：{', '.join(config['focus_points'])}"
                ),
            ]

        contract_type.change(
            fn=update_contract_config,
            inputs=[contract_type],
            outputs=[review_dimensions, focus_points],
        )

        # 文件上传时提取文本
        def extract_text(file):
            if file:
                try:
                    if file.name.endswith(".txt"):
                        with open(
                            file.name, "r", encoding="utf-8", errors="ignore"
                        ) as f:
                            return f.read()
                    elif file.name.endswith(".docx"):
                        doc = Document(file.name)
                        return "\n".join([para.text for para in doc.paragraphs])
                    elif file.name.endswith(".pdf"):
                        reader = PdfReader(file.name)
                        text = ""
                        for page in reader.pages:
                            extracted_text = page.extract_text()
                            if extracted_text:
                                text += extracted_text + "\n"
                        return text
                except Exception as e:
                    return f"文件处理错误: {str(e)}"
            return ""

        file_upload.change(
            fn=extract_text, inputs=[file_upload], outputs=[contract_text]
        )

        # 审核按钮点击事件
        def review_contract(
            contract_type, review_dimensions, focus_points, contract_text
        ):
            if not contract_text.strip():
                return "请输入或上传合同文本后再进行审核"

            result = review_contract_core(
                contract_type, contract_text, review_dimensions, focus_points
            )
            return (
                result["result"]
                if result["success"]
                else f"审核失败: {result['error']}"
            )

        analyze_btn.click(
            fn=review_contract,
            inputs=[contract_type, review_dimensions, focus_points, contract_text],
            outputs=[result],
        )

        # API说明
        gr.Markdown(
            """
        ---
        ## API 接口说明
        可通过API接口集成到您的系统：
        - **POST** `/api/review`
        - **参数**: 合同类型、审核维度、特别关注点、合同文本
        - **返回**: 合同审核结果
        """
        )

    return demo


# --- FastAPI 应用 ---
app = FastAPI(
    title="合同智审 - AI合同审核服务", description="提供专业的合同风险审核与优化建议"
)


# 请求模型
class ReviewRequest(BaseModel):
    contract_type: str
    review_dimensions: List[str]
    focus_points: Optional[str] = ""
    contract_text: str


# 审核接口
@app.post("/api/review", summary="合同审核")
async def api_review(request: ReviewRequest, file: Optional[UploadFile] = File(None)):
    # 如果上传了文件，提取文本并覆盖请求中的文本
    contract_text = request.contract_text
    if file:
        file_text = extract_text_from_file(file)
        if "不支持的文件格式" in file_text or "文件处理错误" in file_text:
            raise HTTPException(status_code=400, detail=file_text)
        contract_text = file_text

    if not contract_text.strip():
        raise HTTPException(status_code=400, detail="请提供合同文本内容")

    result = review_contract_core(
        request.contract_type,
        contract_text,
        request.review_dimensions,
        request.focus_points,
    )

    if result["success"]:
        return result
    else:
        raise HTTPException(status_code=500, detail=result["error"])


# 健康检查
@app.get("/health", summary="健康检查")
async def health_check():
    return {"status": "healthy"}


# 挂载Gradio界面
demo = gradio_interface()
app = gr.mount_gradio_app(app=app, blocks=demo, path="/")

# 启动应用
if __name__ == "__main__":
    print("🚀 启动合同智审服务...")
    print("📱 Gradio 界面: http://localhost:7860")
    print("💻 API 文档: http://localhost:7860/docs")
    uvicorn.run(app, host="0.0.0.0", port=7860)
